from manimlib import *

__curr_dir = os.path.dirname(__file__)
if __curr_dir not in sys.path:
    sys.path.append(__curr_dir)
from MySpring import turns


def Spring3D(start, end, axis):
    """
    创建弹簧，包括首尾节点（球体）和长度标签
    """

    start = np.array(start)
    end = np.array(end)

    global_length = np.linalg.norm(end - start)

    start = axis.coords_to_point(start[0], start[1], start[2])
    end = axis.coords_to_point(end[0], end[1], end[2])

    # 计算起始位置和终点位置之间的向量
    vector = end - start
    # 计算向量的长度
    length = np.linalg.norm(vector)
    # 计算向量的方向
    direction = vector / length

    # 横截面半径
    radius = global_length * 0.05

    curve = mobject.functions.ParametricCurve(
        lambda t: axis.coords_to_point(
            t, radius * np.cos(2 * np.pi * turns * t), radius * np.sin(2 * np.pi * turns * t)  # x坐标  # y坐标  # z坐标
        ),
        t_range=[0, 1, 0.005],  # 绘制步长0.001
        color=YELLOW,
    )
    curve.set_style(stroke_width=10)

    # 将单元弹簧映射到真实坐标上
    curve.shift(start)  # 将弹簧移动到起始位置
    curve.stretch(global_length, 0, about_point=start)  # 根据向量的长度来缩放弹簧
    curve.rotate(np.arccos(direction[0]), about_point=start)  # 根据向量的方向来旋转弹簧

    return curve


def Bar3D(local_start, local_end, axis):
    """
    创建杆
    """

    local_start = np.array(local_start)
    local_end = np.array(local_end)

    global_start = axis.coords_to_point(local_start[0], local_start[1], local_start[2])
    global_end = axis.coords_to_point(local_end[0], local_end[1], local_end[2])

    # 计算起始位置和终点位置之间的向量
    global_vector = global_end - global_start
    # 计算向量的长度
    global_length = np.linalg.norm(global_vector)

    cylinder = Line3D(start=global_start, end=global_end, color=YELLOW, width=global_length * 0.05)

    return cylinder


class Spring2Triangler(Scene):
    def construct(self):
        axes = ThreeDAxes()
        frame = self.camera.frame
        frame.set_height(12)

        # ========================================================
        # ========================================================
        # ========================================================
        # 创建弹簧
        p1 = [0, 0, 0]
        p2 = [5, 0, 0]
        p3 = [3, 3.5, 0]
        self.add(Spring3D(p1, p2, axes), Spring3D(p2, p3, axes), Spring3D(p1, p3, axes))

        # 创建杆
        p1 = [0 + 7, 0, 0]
        p2 = [5 + 7, 0, 0]
        p3 = [3 + 7, 3.5, 0]
        self.add(Bar3D(p1, p2, axes), Bar3D(p2, p3, axes), Bar3D(p1, p3, axes))

        # 创建三角面片
        p1 = np.array([0 + 14, 0, 0])
        p2 = np.array([5 + 14, 0, 0])
        p3 = np.array([3 + 14, 3.5, 0])
        weight1 = 0.6
        weight2 = (1 - weight1) * 0.5
        gauss1 = p1 * weight1 + p2 * weight2 + p3 * weight2
        gauss2 = p1 * weight2 + p2 * weight1 + p3 * weight2
        gauss3 = p1 * weight2 + p2 * weight2 + p3 * weight1
        self.add(
            Polygon(p1, p2, p3).set_fill(YELLOW, opacity=1),
            Dot(gauss1, fill_color=RED),
            Dot(gauss2, fill_color=RED),
            Dot(gauss3, fill_color=RED),
        )
